// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <optional>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::functions::test;

namespace {

class ArrayAverageTest : public FunctionBaseTest {
 protected:
  // Evaluate an expression.
  void testExpr(const VectorPtr& expected, const VectorPtr& input) {
    auto result = evaluate("array_average(C0)", make_row_vector({input}));
    assertEqualVectors(expected, result);
  }
};

} // namespace

// Test floating point arrays.
TEST_F(ArrayAverageTest, doubleInput) {
  auto input = make_nullable_array_vector<double>(
      {{0.0, 1.0, 2.0}, {std::nullopt, 1.0, 2.0}, {std::nullopt}});
  auto expected = make_nullable_flat_vector<double>({1.0, 1.5, std::nullopt});
  testExpr(expected, {input});
}

// Test floating point arrays: limits cases.
TEST_F(ArrayAverageTest, doubleLimitsInput) {
  using limits = std::numeric_limits<double>;
  auto input = make_nullable_array_vector<double>(
      {{limits::infinity(), 1.0, 2.0},
       {limits::lowest(), 1.0, 2.0},
       {limits::max(), 1.0, 2.0},
       {limits::quiet_NaN(), 1.0, 2.0}});
  auto expected = make_nullable_flat_vector<double>(
      {limits::infinity(),
       -5.992310449541053E307,
       5.992310449541053E307,
       limits::quiet_NaN()});
  testExpr(expected, input);
}

// Test empty array in input.
TEST_F(ArrayAverageTest, emptyInput) {
  auto input = make_array_vector<double>({{}});
  auto expected = make_nullable_flat_vector<double>({std::nullopt});
  testExpr(expected, input);
}

// Test all non-null inputs including empty array.
TEST_F(ArrayAverageTest, nonNullInput) {
  auto input = make_array_vector<double>({{}, {1.0, 2.0, 3.0}, {4.0}});
  auto expected = make_nullable_flat_vector<double>({std::nullopt, 2.0, 4.0});
  testExpr(expected, input);
}
